﻿// Learn more about F# at http://fsharp.net
#nowarn "40"
#nowarn "62"

open System
open System.IO
open System.Net
open System.Text
open Microsoft.FSharp.Collections
open Spread

type ProducerState = {
    GroupName   : string
    Channel     : SpreadMessage -> unit
    Leader      : string
    Value       : int
    Subscribers : string Set
    Providers   : string array
}

let isLeader s = s.Leader <> "" && s.Leader = s.GroupName
let noKnownLeader ldr mbrs = not (Seq.exists (fun membr -> membr = ldr) mbrs)

let private handleMembership state change = 
    match change.Reason with
    | Joined newMember when Seq.length change.Members = 1 ->
        { state with Leader = newMember }
    | Joined newMember when isLeader state ->
        let ms = new MemoryStream()
        let bw = new BinaryWriter(ms, Encoding.ASCII)
        state.Value |> bw.Write
        state.Subscribers |> Set.count |> bw.Write 
        state.Subscribers |> Set.iter bw.Write    
        
        Data { 
            Groups = [|newMember|]; 
            ServiceType = MessageFlags.Causal; 
            MessageType = 0us; 
            Data = ms.ToArray() } |> state.Channel
                        
        state
    | _ when noKnownLeader state.Leader change.Members ->
        { state with Leader = Seq.hd change.Members }
    | _ ->
        state

let private publish state subscribers = 
    let pc = Array.length state.Providers
    let idx = Array.IndexOf<string>(state.Providers,state.GroupName)
    let targets = 
        subscribers 
        |> Seq.filter (fun s -> Math.Abs(s.GetHashCode() % pc) = idx)
        |> Seq.to_array
    if targets.Length > 0 then
        let ms = new MemoryStream()
        let bw = new BinaryWriter(ms, Encoding.ASCII)
        state.Value |> bw.Write
        Data { 
            Groups = targets; 
            ServiceType = MessageFlags.Causal; 
            MessageType = 1us; 
            Data = ms.ToArray() } |> state.Channel
    subscribers

let private handleMessage state msg =
    match msg.MessageType with 
    | 0us ->
        let leader = msg.Sender
        let br = new BinaryReader(new MemoryStream(msg.Data), Encoding.ASCII)
        let value = br.ReadInt32()
        let scount = br.ReadInt32()
        let subscribers = Seq.init scount (fun _ -> br.ReadString()) |> Set.of_seq
        printfn "Leader is %s, state is %d (%d subscribers)" leader value scount
        Some { state with Leader = leader; Value = value; Subscribers = subscribers }
    | 1us when state.Leader <> "" ->
        let nextValue = BitConverter.ToInt32(msg.Data, 0)        
        let prefix = if isLeader state then "Leader" else "Replica"
        printfn "%s: %d" prefix nextValue

        let state' = { state with Value = nextValue }
        publish state' state.Subscribers |> ignore
        Some state'
    | 2us ->
        let br = new BinaryReader(new MemoryStream(msg.Data), Encoding.ASCII)
        let subscriber = br.ReadString()
        Some { state with Subscribers = state.Subscribers + publish state (Set.singleton subscriber) }
    | 3us ->
        let br = new BinaryReader(new MemoryStream(msg.Data), Encoding.ASCII)
        let subscriber = br.ReadString()
        Some { state with Subscribers = Set.remove subscriber state.Subscribers }
    | _ -> 
        Some state

let private processEvent state msg = 
    match msg with
    | Message msg ->
        // An application message needs processing, defer to handleMessage
        handleMessage state msg
    | Connection (group, chan) ->
        // We are not connected to the spread daemon
        Console.Title <- group
        Join "Provider" |> chan
        Some { state with GroupName = group; Channel = chan}
    | Transition id ->
        // A membership transition has started. Ignore
        Some state
    | MembershipChange change ->
        // A membership changed has occurred
        Some { (handleMembership state change) with Providers = Array.of_seq change.Members }
    | _ ->
        // Otherwise, we are done.
        None

// Spawn our agent, passing the endPoint of the spread daemon
let private startAgent endPoint = 
    let events = MailboxProcessor.Start(fun inbox ->
        let rec loop state = 
            async {
                // The leader updates state every second, replicas just wait.
                let timeout = if isLeader state then 1000 else System.Threading.Timeout.Infinite
                let! result = inbox.TryReceive(timeout)
                match result with
                | Some msg ->
                    // An inbound message needs processing
                    match processEvent state msg with
                    | Some state' ->
                        do! loop state'
                    | None -> ()
                | None ->
                    // Timeout waiting for message. Update state.
                    Data { 
                        Groups = [|"Provider"|]; 
                        ServiceType = MessageFlags.Causal; 
                        MessageType = 1us; 
                        Data = BitConverter.GetBytes(state.Value + 1) } |> state.Channel
                    do! loop state
            }
        loop { GroupName = ""; Channel = ignore; Leader = ""; Value = 0; Subscribers = Set.empty; Providers = Array.create 0 String.Empty })

    // Now that we have defined the mailbox, start the spread client and return
    // the mailbox
    events.Post |> Spread.runClient endPoint true None
    events

match System.Environment.GetCommandLineArgs() with
| [|_;ipaddr;port|] ->
    let ep = IPEndPoint(IPAddress.Parse(ipaddr), Int32.of_string(port))
    let mbox = startAgent ep
    System.Console.ReadLine() |> ignore
| _ ->
    printfn "Usage: FSpreadProvider <idaddr> <port>"